﻿namespace WgFrame.Security
{
    using System;
    using System.Security.Cryptography;
    using UnityEngine;

    internal class DiffieHellmanCryptoProvider : ICryptoProvider, IDisposable
    {
        private readonly BigInteger primeRoot = new BigInteger((long)OakleyGroups.Generator);
        private readonly BigInteger prime = new BigInteger(OakleyGroups.OakleyPrime768);
        private readonly BigInteger secret;
        private readonly BigInteger publicKey;
        private Rijndael crypto;
        private byte[] sharedKey;

        public DiffieHellmanCryptoProvider()
        {
            secret = GenerateRandomSecret(160);
            publicKey = CalculatePublicKey();
        }

        public DiffieHellmanCryptoProvider(byte[] cryptoKey)
        {
            crypto = new RijndaelManaged();
            crypto.Key = cryptoKey;
            crypto.IV = new byte[0x10];
            crypto.Padding = PaddingMode.PKCS7;
        }

        private BigInteger CalculatePublicKey()
        {
            return primeRoot.modPow(secret, prime);
        }

        private BigInteger CalculateSharedKey(BigInteger otherPartyPublicKey)
        {
            return otherPartyPublicKey.modPow(secret, prime);
        }

        public byte[] Decrypt(byte[] data)
        {
            return Decrypt(data, 0, data.Length);
        }
            

        public byte[] Decrypt(byte[] data, int offset, int count)
        {
            using (ICryptoTransform transform = crypto.CreateDecryptor())
            {
                return transform.TransformFinalBlock(data, offset, count);
            }
        }

        public void DeriveSharedKey(byte[] otherPartyPublicKey)
        {
            byte[] buffer;
            BigInteger integer = new BigInteger(otherPartyPublicKey);
            sharedKey = CalculateSharedKey(integer).getBytes();
            using (SHA256 sha = new SHA256Managed())
            {
                buffer = sha.ComputeHash(sharedKey);
            }

            crypto = new RijndaelManaged();
            crypto.Key = buffer;
            crypto.IV = new byte[0x10];
            crypto.Padding = PaddingMode.PKCS7;
            crypto.Mode = CipherMode.ECB;
        }

        public void Dispose()
        {
            Dispose(true);
            GC.SuppressFinalize(this);
        }

        protected void Dispose(bool disposing)
        {
            if (!disposing)
            {
            }
        }

        public byte[] Encrypt(byte[] data)
        {
            return Encrypt(data, 0, data.Length);
        }
            

        public byte[] Encrypt(byte[] data, int offset, int count)
        {
            using (ICryptoTransform transform = crypto.CreateEncryptor())
            {
                return transform.TransformFinalBlock(data, offset, count);
            }
        }

        public BigInteger GenerateRandomSecret(int secretLength)
        {
            while (true)
            {
                BigInteger integer = new BigInteger();
                integer.genRandomBits(secretLength, new RNGCryptoServiceProvider());
                if ((integer < (prime - 1)) && (integer != 0))
                {
                    return integer;
                }
            }
        }

        public bool IsInitialized
        { 
            get
            {
                return (crypto != null);
            }
        }
            

        public byte[] PublicKey
        { 
            get
            {
                return publicKey.getBytes();
            }
        }

    }
}

